// ask.js
const fs = require("fs");
const path = require("path");
const { RecursiveCharacterTextSplitter } = require("langchain/text_splitter");
const { RetrievalQAChain } = require("langchain/chains");
const { TextLoader } = require("langchain/document_loaders/fs/text");
const { Chroma } = require("@langchain/community/vectorstores/chroma");
const { ChatOpenAI } = require("@langchain/openai");

// 简单的本地向量化实现
class SimpleEmbeddings {
  constructor() {
    // 固定维度，避免复杂计算
    this.dimension = 1536; // 使用标准维度
  }

  async embedDocuments(documents) {
    const embeddings = [];
    for (const doc of documents) {
      const content = doc.page_content || doc;
      const vector = this.textToVector(content);
      embeddings.push(vector);
    }
    return embeddings;
  }

  async embedQuery(text) {
    return this.textToVector(text);
  }

  textToVector(text) {
    // 简单的哈希向量化方法，生成固定长度的向量
    const vector = new Array(this.dimension).fill(0);
    const normalizedText = text.toLowerCase().replace(/[^\w\s]/g, "");

    // 使用简单的哈希方法分配权重
    for (let i = 0; i < Math.min(text.length, 100); i++) {
      const charCode = text.charCodeAt(i);
      const index = (charCode * 31 + i) % this.dimension;
      vector[index] += Math.random();
    }

    // 归一化向量
    const magnitude = Math.sqrt(
      vector.reduce((sum, val) => sum + val * val, 0)
    );
    if (magnitude > 0) {
      return vector.map((val) => val / magnitude);
    }
    return vector;
  }

  getDimension() {
    return this.dimension;
  }
}

class LocalQASystem {
  constructor() {
    this.chain = null;
    this.embeddings = null;
    this.vectordb = null;
    this.dataDir = path.join(__dirname, "data");
  }

  async initSystem() {
    try {
      console.log("正在初始化问答系统...");
      const startTime = Date.now();

      // 1. 处理文档
      console.log("1. 处理文档...");
      const documents = await this.processDocuments();

      // 2. 创建向量数据库
      console.log("2. 创建向量数据库...");
      await this.createVectorDB(documents);

      // 3. 创建问答链
      console.log("3. 创建问答链...");
      await this.createQAChain();

      const endTime = Date.now();
      console.log(`系统初始化完成，耗时: ${(endTime - startTime) / 1000}秒`);
      return true;
    } catch (error) {
      console.error(`系统初始化失败: ${error.message}`);
      console.error(error.stack);
      return false;
    }
  }

  async processDocuments() {
    try {
      // 确保data目录存在
      if (!fs.existsSync(this.dataDir)) {
        fs.mkdirSync(this.dataDir, { recursive: true });
        console.log(`警告: ${this.dataDir} 目录为空，请添加txt文档`);
        return [];
      }

      // 加载文档
      const documents = await this.loadDocuments(this.dataDir);
      if (documents.length === 0) {
        console.log("警告: 没有找到任何文档");
        return [];
      }

      console.log(`  共加载 ${documents.length} 个文档`);
      return documents;
    } catch (error) {
      console.error(`处理文档时出错: ${error.message}`);
      return [];
    }
  }

  async loadDocuments(dataDir) {
    try {
      console.log(`正在从目录加载文档: ${dataDir}`);

      // 查找所有txt文件
      const txtFiles = [];
      const files = fs.readdirSync(dataDir);
      console.log(`在目录中找到 ${files.length} 个文件`);

      for (const file of files) {
        if (file.endsWith(".txt") && !file.startsWith("cut_")) {
          const fullPath = path.join(dataDir, file);
          txtFiles.push(fullPath);
          console.log(`  发现txt文件: ${file}`);
        }
      }

      if (txtFiles.length === 0) {
        console.log("警告: 没有找到txt文档");
        return [];
      }

      const documents = [];
      for (const filePath of txtFiles) {
        try {
          console.log(`  正在加载: ${filePath}`);
          const loader = new TextLoader(filePath);
          const docs = await loader.load();
          documents.push(...docs);
          console.log(`    成功加载 ${docs.length} 个文档片段`);
        } catch (error) {
          console.error(`  加载 ${filePath} 时出错: ${error.message}`);
        }
      }

      console.log(`总共加载了 ${documents.length} 个文档`);
      return documents;
    } catch (error) {
      console.error(`加载文档时出错: ${error.message}`);
      return [];
    }
  }

  async createVectorDB(documents) {
    try {
      if (documents.length === 0) {
        console.log("警告: 没有文档可处理");
        return null;
      }

      // 文档分割
      const textSplitter = new RecursiveCharacterTextSplitter({
        chunkSize: 800,
        chunkOverlap: 100,
      });

      const splitDocs = await textSplitter.splitDocuments(documents);
      console.log(`  文档分割完成，共 ${splitDocs.length} 个片段`);

      // 创建嵌入模型
      console.log("正在加载Embeddings模型...");
      this.embeddings = new SimpleEmbeddings();
      console.log(
        `Embeddings模型加载完成，维度: ${this.embeddings.getDimension()}`
      );

      // 创建向量数据库
      console.log("正在创建向量数据库...");

      // 先尝试清理已存在的集合
      try {
        const { ChromaClient } = require("chromadb");
        const chromaClient = new ChromaClient({
          host: "localhost",
          port: 8000,
        });

        // 检查集合是否存在并删除
        try {
          await chromaClient.deleteCollection({
            name: "local_qa_collection",
          });
          console.log("已删除已存在的集合");
        } catch (error) {
          // 集合可能不存在，这是正常的
          console.log("集合不存在或删除失败（可能是首次创建）:", error.message);
        }
      } catch (error) {
        console.log("Chroma客户端初始化失败:", error.message);
      }

      // 使用新的配置方式创建向量数据库
      this.vectordb = await Chroma.fromDocuments(splitDocs, this.embeddings, {
        collectionName: "local_qa_collection",
        client: {
          host: "localhost",
          port: 8000,
        },
      });

      console.log(`向量数据库创建完成，共 ${splitDocs.length} 个文档片段`);
      return this.vectordb;
    } catch (error) {
      console.error(`创建向量数据库时出错: ${error.message}`);
      console.error("错误堆栈:", error.stack);
      return null;
    }
  }

  async createQAChain() {
    try {
      if (this.vectordb === null) {
        console.log("错误: 向量数据库未初始化");
        return null;
      }

      console.log("正在创建问答链...");

      // 创建DeepSeek聊天模型
      const llm = new ChatOpenAI({
        temperature: 0.1,
        modelName: "deepseek-chat",
        openAIApiKey: "sk-980ecc7694e14620a29b5d679a3d1c05",
        configuration: {
          baseURL: "https://api.deepseek.com/v1",
        },
      });

      // 创建检索器
      const retriever = this.vectordb.asRetriever({
        searchType: "similarity",
        k: 4,
      });

      // 创建问答链
      this.chain = new RetrievalQAChain({
        llm: llm,
        retriever: retriever,
        returnSourceDocuments: true,
      });

      console.log("问答链创建完成");
      return this.chain;
    } catch (error) {
      console.error(`创建问答链时出错: ${error.message}`);
      return null;
    }
  }

  async askQuestion(question) {
    if (this.chain === null) {
      return "系统未初始化，请先初始化系统。";
    }

    try {
      console.log(`正在思考问题: ${question}`);
      const startTime = Date.now();

      // 调用问答链
      const result = await this.chain.invoke({ query: question });

      const endTime = Date.now();
      console.log(`回答完成，耗时: ${(endTime - startTime) / 1000}秒`);

      // 检查是否有相关文档
      if (result.sourceDocuments && result.sourceDocuments.length > 0) {
        return result.text;
      } else {
        return "抱歉，我在知识库中没有找到相关信息来回答您的问题。";
      }
    } catch (error) {
      console.error("问答错误:", error);
      return `回答问题时出现错误: ${error.message}`;
    }
  }

  getSystemStatus() {
    if (this.vectordb === null) {
      return { status: "未初始化", document_count: 0 };
    }

    try {
      return { status: "已初始化", document_count: "未知" };
    } catch (error) {
      return { status: "已初始化", document_count: "未知" };
    }
  }

  async getDatabaseInfo(page = 1, perPage = 5) {
    if (this.vectordb === null) {
      return { error: "数据库未初始化" };
    }

    try {
      return {
        document_count: "未知",
        current_page: page,
        total_pages: 1,
        per_page: perPage,
        document_previews: [],
      };
    } catch (error) {
      return { error: error.message };
    }
  }
}

module.exports = { LocalQASystem };
